package org.easy.byteasr;


import com.alibaba.fastjson2.JSONObject;
import lombok.extern.slf4j.Slf4j;
import okhttp3.OkHttpClient;
import okhttp3.Request;
import okhttp3.Response;
import okhttp3.WebSocket;
import okhttp3.WebSocketListener;
import okhttp3.logging.HttpLoggingInterceptor;
import okhttp3.logging.HttpLoggingInterceptor.Level;
import okio.ByteString;
import org.easy.byteasr.dto.AudioFormat;
import org.easy.byteasr.dto.ByteAudioInfoResponse;

import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.util.HashMap;
import java.util.Map;
import java.util.UUID;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.TimeUnit;
import java.util.zip.GZIPInputStream;
import java.util.zip.GZIPOutputStream;

/**
 * 单包音频必须在100ms大小左右
 * https://www.volcengine.com/docs/6561/1354869
 */
@Slf4j
public class ByteRealtimeAsrV3Client {

    private static final byte PROTOCOL_VERSION = 0b0001;
    private static final byte DEFAULT_HEADER_SIZE = 0b0001;
    private static final byte FULL_CLIENT_REQUEST = 0b0001;
    private static final byte AUDIO_ONLY_REQUEST = 0b0010;
    private static final byte FULL_SERVER_RESPONSE = 0b1001;
    private static final byte SERVER_ACK = 0b1011;
    private static final byte SERVER_ERROR_RESPONSE = 0b1111;
    // Message Type Specific Flags
    private static final byte NO_SEQUENCE = 0b0000;// no check sequence
    private static final byte POS_SEQUENCE = 0b0001;
    private static final byte NEG_SEQUENCE = 0b0010;
    private static final byte NEG_WITH_SEQUENCE = 0b0011;
    private static final byte NEG_SEQUENCE_1 = 0b0011;
    // Message Serialization
    private static final byte NO_SERIALIZATION = 0b0000;
    private static final byte JSON = 0b0001;
    // Message Compression
    private static final byte NO_COMPRESSION = 0b0000;
    private static final byte GZIP = 0b0001;
    private final String url = "wss://openspeech.bytedance.com/api/v3/sauc/bigmodel";
    private final String appId;
    private final String token;
    public ByteRealtimeAsrV3Client(String appId, String token) {
        this.appId = appId;
        this.token = token;
    }

    static byte[] getHeader(byte messageType, byte messageTypeSpecificFlags, byte serialMethod,
                            byte compressionType, byte reservedData) {
        final byte[] header = new byte[4];
        header[0] = (PROTOCOL_VERSION << 4) | DEFAULT_HEADER_SIZE; // Protocol version|header size
        header[1] = (byte) ((messageType << 4) | messageTypeSpecificFlags); // message type | messageTypeSpecificFlags
        header[2] = (byte) ((serialMethod << 4) | compressionType);
        header[3] = reservedData;
        return header;
    }

    static byte[] intToBytes(int a) {
        return new byte[]{
                (byte) ((a >> 24) & 0xFF),
                (byte) ((a >> 16) & 0xFF),
                (byte) ((a >> 8) & 0xFF),
                (byte) (a & 0xFF)

        };
    }

    static int bytesToInt(byte[] src) {
        if (src == null || (src.length != 4)) {
            throw new IllegalArgumentException("");
        }
        return ((src[0] & 0xFF) << 24)
                | ((src[1] & 0xff) << 16)
                | ((src[2] & 0xff) << 8)
                | ((src[3] & 0xff));
    }

    static byte[] generateBeforePayload(int seq) {
        return intToBytes(seq);
    }

    static byte[] gzipCompress(byte[] src) {
        return gzipCompress(src, src.length);
    }

    static byte[] gzipCompress(byte[] src, int len) {
        if (src == null || len == 0) {
            return new byte[0];
        }
        ByteArrayOutputStream out = new ByteArrayOutputStream();
        GZIPOutputStream gzip = null;
        try {
            gzip = new GZIPOutputStream(out);
            gzip.write(src, 0, len);
        } catch (IOException e) {
            e.printStackTrace();
        } finally {
            if (gzip != null) {
                try {
                    gzip.close();
                } catch (IOException e) {
                    e.printStackTrace();
                }
            }
        }
        return out.toByteArray();
    }

    static byte[] gzipDecompress(byte[] src) {
        if (src == null || src.length == 0) {
            return null;
        }
        ByteArrayOutputStream out = new ByteArrayOutputStream();
        ByteArrayInputStream ins = new ByteArrayInputStream(src);
        GZIPInputStream gzip = null;
        try {
            gzip = new GZIPInputStream(ins);
            byte[] buffer = new byte[ins.available()];
            int len = 0;
            while ((len = gzip.read(buffer)) > 0) {
                out.write(buffer, 0, len);
            }
            out.close();
        } catch (IOException e) {
            e.printStackTrace();
        } finally {
            if (gzip != null) {
                try {
                    gzip.close();
                } catch (IOException e) {
                    e.printStackTrace();
                }
            }
        }

        return out.toByteArray();
    }

    static ByteAudioInfoResponse parserResponse(byte[] res) {
        if (res == null || res.length == 0) {
            return null;
        }
        // 当符号位为1时进行 >> 运算后高位补1（预期是补0），导致结果错误，所以增加个数再与其& 运算，目的是确保高位是补0.
        final byte num = 0b00001111;
        Map<String, Object> result = new HashMap<>();
        // header 32 bit=4 byte
        int protocol_version = (res[0] >> 4) & num;
        result.put("protocol_version", protocol_version);
        int header_size = res[0] & 0x0f;
        result.put("header_size", header_size);

        int message_type = (res[1] >> 4) & num;
        result.put("message_type", message_type);
        int message_type_specific_flags = res[1] & 0x0f;
        result.put("message_type_specific_flags", message_type_specific_flags);
        int serialization_method = res[2] >> num;
        result.put("serialization_method", serialization_method);
        int message_compression = res[2] & 0x0f;
        result.put("message_compression", message_compression);
        int reserved = res[3];
        result.put("reserved", reserved);

        // sequence 4 byte
        byte[] temp = new byte[4];
        System.arraycopy(res, 4, temp, 0, temp.length);
        int sequence = bytesToInt(temp);// sequence 4 byte


        // payload size 4 byte
        String payloadStr = null;
        System.arraycopy(res, 8, temp, 0, temp.length);
        int payloadSize = bytesToInt(temp);
        byte[] payload = new byte[res.length - 12];
        System.arraycopy(res, 12, payload, 0, payload.length);
        // 正常Response
        if (message_type == FULL_SERVER_RESPONSE) {
            if (message_compression == GZIP) {
                payloadStr = new String(gzipDecompress(payload));
            } else {
                payloadStr = new String(payload);
            }
            result.put("payload_size", payloadSize);
            return com.alibaba.fastjson2.JSON.parseObject(payloadStr, ByteAudioInfoResponse.class);
        } else if (message_type == SERVER_ACK) {
            payloadStr = new String(payload);
            result.put("payload_size", payloadSize);
            return com.alibaba.fastjson2.JSON.parseObject(payloadStr, ByteAudioInfoResponse.class);
        } else if (message_type == SERVER_ERROR_RESPONSE) {
            payloadStr = new String(payload);
            result.put("code", sequence);
            result.put("error msg", payloadStr);
            System.out.println("===>response:" + JSONObject.toJSONString(result));
            return com.alibaba.fastjson2.JSON.parseObject(payloadStr, ByteAudioInfoResponse.class);
        }
        return null;

    }

    public ByteAsrWebsocketHolder init(ByteAsrWebSocketListener webSocketListener, CountDownLatch latch, AudioFormat audioFormat) {

        final Request request = new Request.Builder()
                .url(url)
                .header("X-Api-App-Key", appId)
                .header("X-Api-Access-Key", token)
                .header("X-Api-Resource-Id", "volc.bigasr.sauc.duration")
                .header("X-Api-Connect-Id", UUID.randomUUID().toString())
                .build();

        HttpLoggingInterceptor interceptor = new HttpLoggingInterceptor();
        interceptor.setLevel(Level.HEADERS);
        final OkHttpClient okHttpClient = new OkHttpClient.Builder().pingInterval(50, TimeUnit.SECONDS)
                .addInterceptor(interceptor)
                .readTimeout(100, TimeUnit.SECONDS)
                .writeTimeout(100, TimeUnit.SECONDS)
                .build();


        final String sessionId = UUID.randomUUID().toString().replace("-", "");
        WebSocket webSocket = okHttpClient.newWebSocket(request, new WebSocketListener() {

            @Override
            public void onOpen(WebSocket webSocket, Response response) {
                System.out.println("===> onOpen,X-Tt-Logid:" + response.header("X-Tt-Logid"));
                webSocketListener.onOpen(webSocket, response);
                latch.countDown();
            }

            @Override
            public void onMessage(WebSocket webSocket, ByteString bytes) {
                byte[] res = bytes.toByteArray();
                ByteAudioInfoResponse byteAudioInfoResponse = parserResponse(res);
                webSocketListener.onMessage(webSocket, byteAudioInfoResponse);
            }

            @Override
            public void onMessage(WebSocket webSocket, String text) {
                super.onMessage(webSocket, text);
            }


            @Override
            public void onClosing(WebSocket webSocket, int code, String reason) {
                super.onClosing(webSocket, code, reason);
                System.out.println("===> onClosing： code:" + code + " reason:" + reason);

            }

            @Override
            public void onClosed(WebSocket webSocket, int code, String reason) {
                super.onClosed(webSocket, code, reason);
                System.out.println("===> onClosed： code:" + code + " reason:" + reason);
                webSocketListener.onClose(webSocket, code, reason);
            }

            @Override
            public void onFailure(WebSocket webSocket, Throwable t, Response response) {
                super.onFailure(webSocket, t, response);
                System.out.println("===> onFailure： Throwable:" + t.getMessage() + " Response:" + (response == null ? "null" : response.toString()));
                webSocketListener.onClose(webSocket,-1 , t.getMessage());
            }
        });

        ByteAsrWebsocketHolder holder = new ByteAsrWebsocketHolder();
        holder.setSessionId(sessionId);
        holder.setWebSocket(webSocket);
        holder.setSeq(1);
        holder.setAudioFormat(audioFormat);
        holder.setListener(webSocketListener);
        return holder;

    }

    public boolean sendAudioFirst(WebSocket webSocket, String userId, AudioFormat format) {
        byte[] header = getHeader(FULL_CLIENT_REQUEST, POS_SEQUENCE, JSON, GZIP, (byte) 0);
        byte[] seqBytes = generateBeforePayload(1);
        final byte[] payloadBytes = gzipCompress(getPlayload(userId, format).getBytes());
        final byte[] payloadSize = intToBytes(payloadBytes.length);
        final byte[] fullClientRequest = buildRequest(header, seqBytes, payloadSize, payloadBytes);
        return webSocket.send(ByteString.of(fullClientRequest));

    }


    private String getPlayload(String uid, AudioFormat format) {
        JSONObject user = new JSONObject();
        user.put("uid", uid);

        JSONObject audio = new JSONObject();
        audio.put("format", "pcm");
        audio.put("sample_rate", (int) format.getSampleRate());
        audio.put("bits", format.getSampleSizeInBits());
        audio.put("channel", format.getChannels());
        audio.put("codec", "raw");

        JSONObject request = new JSONObject();
        request.put("model_name", "bigmodel");
        request.put("enable_punc", false);
        request.put("show_utterances", true);
        request.put("result_type", "single");

        JSONObject payload = new JSONObject();
        payload.put("user", user);
        payload.put("audio", audio);
        payload.put("request", request);

        return payload.toString();
    }

    // audio_only_request= header + sequence + payload size+ payload
    public boolean sendAudioOnlyWithSeq(WebSocket webSocket, byte[] buffer, int len, boolean isLast, int seq) {

        byte messageTypeSpecificFlags = isLast ? NEG_WITH_SEQUENCE : POS_SEQUENCE;
        // header
        byte[] header = getHeader(AUDIO_ONLY_REQUEST, messageTypeSpecificFlags, JSON, GZIP, (byte) 0);
        // sequence
        byte[] sequenceBytes = generateBeforePayload(seq);
        // payload size
        byte[] payloadBytes = gzipCompress(buffer, len);
        // payload
        byte[] payloadSize = intToBytes(payloadBytes.length);
        byte[] audio_only_request = buildRequest(header, sequenceBytes, payloadSize, payloadBytes);
        return webSocket.send(ByteString.of(audio_only_request));
    }

    private byte[] buildRequest(byte[] header, byte[] seqBytes, byte[] payloadSize, byte[] payloadBytes) {
        byte[] request = new byte[header.length + seqBytes.length + payloadSize.length + payloadBytes.length];
        int destPos = 0;
        System.arraycopy(header, 0, request, destPos, header.length);
        destPos += header.length;
        System.arraycopy(seqBytes, 0, request, destPos, seqBytes.length);
        destPos += seqBytes.length;
        System.arraycopy(payloadSize, 0, request, destPos, payloadSize.length);
        destPos += payloadSize.length;
        System.arraycopy(payloadBytes, 0, request, destPos, payloadBytes.length);
        return request;
    }

}
